# Autor chenfeng
# !/usr/bin/env Python
# coding=utf-8


from paddlenlp import Taskflow


#########################10、主题提取
# todo [模型下载有问题！！！！]
# summary = Taskflow("text_summarization")
# print('#########################10、主题提取')
# # print(summary('2021年10月，PaddleNLP 聚合众多百度自然语言处理领域自研算法以及社区开源模型，并凭借飞桨核心框架的能力升级开放了开箱即用、极致优化的高性能一键预测功能，备受开发者喜爱。'))
# sum = summary("2022年，中国房地产进入转型阵痛期，传统“高杠杆、快周转”的模式难以为继，万科甚至直接喊话，中国房地产进入“黑铁时代” ")
# print(sum)

#########################11、代码生成
# todo [模型下载有问题！！！！] 参考 https://cloud.tencent.com/developer/article/2178802
#Taskflow调用
# codegen = Taskflow("code_generation", model="Salesforce/codegen-2B-mono",decode_strategy="greedy_search", repetition_penalty=1.0)
# prompt = "def twoSum(nums, target):\n    hashmap={}\n    for ind,num in enumerate(nums):\n        hashmap[num] = ind\n    for i,num in enumerate(nums):"
# print("#########################11、代码补全")
# code = codegen(prompt)
# print('原文：\n' + prompt)
# print('代码续写：\n' + code[0])


#########################12、文档智能（文档抽取问答）
# todo 参考 https://aistudio.baidu.com/aistudio/projectdetail/4881278?channelType=0&channel=0
docprompt = Taskflow("document_intelligence", model='docprompt')
print("#########################13、文档智能（文档抽取问答）")
'''
[{'prompt': '发票金额', 'result': [{'value': '97. 09', 'prob': 0.99, 'start': 134, 'end': 137}]}, 
 {'prompt': '右侧抬头日期是什么?', 'result': [{'value': '2021 年 04月01日', 'prob': 0.76, 'start': 31, 'end': 38}]}, 
 {'prompt': '纳税人识别号是什么？', 'result': [{'value': '91440300MA5FYWC33A', 'prob': 1.0, 'start': 76, 'end': 76}]}]
'''
result = docprompt({"doc": "./invoice.jpg", "prompt": ["发票金额", "右侧抬头日期是什么?", "纳税人识别号是什么？"]})
print(result)


#########################13、文字转图片
print("#########################12、文字转图片")
def text_to_image2(prompt):
    from paddlenlp.transformers.dallebart.tokenizer import DalleBartTokenizer
    from paddlenlp.transformers.dallebart.modeling import DalleBartForImageGeneration
    # todo 仿照Taskflow自定义的image_taskflow
    #   参考 https://aistudio.baidu.com/aistudio/projectdetail/4413005?channelType=0&channel=0
    from image_taskflow import Taskflow as image_taskflow

    # 加载模型和tokenizer
    model_name_or_path = "dalle-mini" #dalle-mega
    model = DalleBartForImageGeneration.from_pretrained(model_name_or_path)
    tokenizer = DalleBartTokenizer.from_pretrained(model_name_or_path)
    # 创建文字转换为图像实例
    text_to_image = image_taskflow("text_to_image", tokenizer, model, tokenizer,num_return_sequences=8)
    text_to_image(prompt)


def text_to_image1(prompt):
    # todo 参考 https://paddlenlp.readthedocs.io/zh/latest/_modules/paddlenlp/taskflow/text_to_image.html#TextToImageGenerationTask
    # 默认模型准确性不是很好
    text_to_image = Taskflow('text_to_image')
    # todo dalle-mega模型有9.6G （虽然更准确，但太大了，没有下载）
    # text_to_image = Taskflow('text_to_image', model = "dalle-mini")

    print(prompt)
    images = text_to_image(prompt)
    for images_ in images:
        for image_ in images_:
            image_.show()


# prompt = "奥特骑车"
# prompt = "风阁水帘今在眼，且来先看早梅红"
# prompt = "A big cat is eating a robot cow"
# prompt = "A little prince with a rose in his hand and a fox were sitting alone on a star looking out into the distance" # 小王子插图
prompts = [
    "The cartoon head of Rick and Morty",   # 瑞克和莫蒂的卡通头像
    "an armchair in the shape of an avocado", # 一个牛油果形状的单人沙发
    "Fireworks are being set off at the Bird's Nest Stadium", # 鸟巢体育场正在燃放烟花
    "A Van Gogh-style shoulder bag", # 一个具有梵高风格的单肩包
]
# 输出
for prompt in prompts:
    print(prompt)
    text_to_image2(prompt)


#########################14、提问生成
print("#########################14、提问生成")
que = Taskflow("question_generation")
'''
question =  ['黄山最高峰']
'''
# question = que([{"context": "奇峰黄山千米以上的山峰有77座，整座黄山就是一座花岗岩的峰林，自古有36大峰，36小峰，" \
#                         "最高峰莲花峰、最险峰天都峰和观日出的最佳点光明顶构成黄山的三大主峰。", "answer": "莲花峰"}])
'''
question = ['苹果手机不见了用什么唤醒软件']
'''
# question = que([{"context": "我的手机不见了，我尝试用语音助手唤醒功能找到了它", "answer": "语音助手唤醒功能"}])
'''
question = ['苹果手机不见了用什么软件']
'''
question = que([{"context": "我的手机不见了", "answer": "语音助手唤醒功能试下"}])

print(question)
